#include "StructureTensor.h"

#include "ST_EstimateFunction.h"

#include <stdio.h>
#include <float.h>
#include <memory.h>
#include <math.h>


void FiberParameterization(CFiber *fiber, CFiber *altas, CPoint3F *parameters, CPoint3F &x, const int stepSize)
{
	/* first check if we need reverse the fiber */
	float reverseFlag = true;
	if (InnerProduct(fiber->m_seed.dir, altas->m_seed.dir) > 0.0f) {
		reverseFlag = false;
	} else {
		reverseFlag = true;
	}

    /* resample the fiber */
    int size = fiber->m_fCount + fiber->m_bCount - 2;
    int resampleSize = 0;
    CPoint3F *resample = NULL;
    if (size > 1) {
        resampleSize = size * RESAMPLE + 1;
        resample = new CPoint3F[resampleSize];
        int k = 0;
        for (int i = fiber->m_bCount-1; i > 0; --i) {
            CVector3F v = fiber->m_pB[i-1] - fiber->m_pB[i];
            float len = v.Length();
            float step = len / (float)RESAMPLE;
            
           /*printf("(%f %f %f), (%f %f %f), v = (%f %f %F), len = %f, step = %f\n", 
                    fiber->m_pB[i-1].m_x, fiber->m_pB[i-1].m_y, fiber->m_pB[i-1].m_z, 
                    fiber->m_pB[i].m_x, fiber->m_pB[i].m_y, fiber->m_pB[i].m_z, 
                    v.m_x, v.m_y, v.m_z, len, step);*/
            
            //v.Normalize();
            for (int j = 0; j < RESAMPLE; ++j) {
                resample[k] = fiber->m_pB[i] + v * (step * (float)j);
                k++;
            }
        }
        resample[k] = fiber->m_pB[0];
        k++;
        for (int i = 1; i < fiber->m_fCount; ++i) {
            CVector3F v = fiber->m_pF[i] - fiber->m_pF[i-1];
            float len = v.Length();
            float step = len / (float)RESAMPLE;

            /*printf("(%f %f %f), (%f %f %f), v = (%f %f %F), len = %f, step = %f\n", 
                    fiber->m_pF[i-1].m_x, fiber->m_pF[i-1].m_y, fiber->m_pF[i-1].m_z, 
                    fiber->m_pF[i].m_x, fiber->m_pF[i].m_y, fiber->m_pF[i].m_z, 
                    v.m_x, v.m_y, v.m_z, len, step);*/

            //v.Normalize();
            for (int j = 1; j <= RESAMPLE; ++j) {
                resample[k] = fiber->m_pF[i-1] + v * (step * (float)j);
                k++;
            }
        }
    } else {
        resampleSize = 1;
        resample = new CPoint3F[resampleSize];
        resample[0] = fiber->m_pF[0];
    }

	
	/* then compute the closest point on fiber to the center, and set it as x */
	float distance = FLT_MAX;
	x = fiber->m_seed.pos;
    int id = -1;
	for (int i = 0; i < resampleSize; ++i) {
		//float temp = Distance(resample[i], altas->m_seed.pos);
        float temp = Distance(resample[i], fiber->m_seed.pos);
		if (temp < distance) {
			distance = temp;
			x = resample[i];
            id = i;
		}
	}
	
	
	/* finally, do the parameterization */
	if (reverseFlag) {
        parameters[PARAM_COUNT/2] = resample[id];
        for (int i = 1; i <= PARAM_COUNT/2; ++i) {
            int k = i * stepSize;
            int index = (id + k) < resampleSize ? (id + k) : (resampleSize - 1);
            parameters[PARAM_COUNT/2+i] = resample[index];
            index = (id - k) > 0 ? (id - k) : 0;
            parameters[PARAM_COUNT/2-i] = resample[index];
        }
	} else {
        parameters[PARAM_COUNT/2] = resample[id];
        for (int i = 1; i <= PARAM_COUNT/2; ++i) {
            int k = i * stepSize;
            int index = (id + k) < resampleSize ? (id + k) : (resampleSize - 1);
            parameters[PARAM_COUNT/2-i] = resample[index];
            index = (id - k) > 0 ? (id - k) : 0;
            parameters[PARAM_COUNT/2+i] = resample[index];
        }
	}

    /* debug, print the fiber and parameters */
    /*printf("fiber:\n");
    for (int i = 0; i < fiber->m_fCount; ++i) {
        printf("(%f %f %f)\n", fiber->m_pF[i].m_x, fiber->m_pF[i].m_y, fiber->m_pF[i].m_z);
    }
    for (int i = 0; i < fiber->m_bCount; ++i) {
        printf("(%f %f %f)\n", fiber->m_pB[i].m_x, fiber->m_pB[i].m_y, fiber->m_pB[i].m_z);
    }
    printf("resample:\n");
    for (int i = 0; i < resampleSize; ++i) {
        printf("(%f %f %f)\n", resample[i].m_x, resample[i].m_y, resample[i].m_z);
    }
    printf("parameters:\n");
    for (int i = 0; i < PARAM_COUNT; ++i) {
        printf("(%f %f %f)\n", parameters[i].m_x, parameters[i].m_y, parameters[i].m_z);
    }
    printf("id = %d, x = (%f %f %f)\n", id, x.m_x, x.m_y, x.m_z);
    printf("\n");*/

    delete[] resample;
}


/* compute the edge strength from the estimated function */
void  ComputeStructureTensor(float *function, float *st)
{
    float *gX = new float[PARAM_COUNT*3];
    float *gY = new float[PARAM_COUNT*3];
    float *gZ = new float[PARAM_COUNT*3];

    for (int i = 0; i < PARAM_COUNT; ++i) {
        gX[i*3+0] = function[i*9+0];
        gX[i*3+1] = function[i*9+3];
        gX[i*3+2] = function[i*9+6];
        gY[i*3+0] = function[i*9+1];
        gY[i*3+1] = function[i*9+4];
        gY[i*3+2] = function[i*9+7];
        gZ[i*3+0] = function[i*9+3];
        gZ[i*3+1] = function[i*9+6];
        gZ[i*3+2] = function[i*9+8];
    }

	st[0] = 1.0f;
	st[1] = 0.0f;
	st[2] = 0.0f;
	st[3] = 0.0f;
	st[4] = 0.0f;
	st[5] = 0.0f;
	st[6] = 0.0f;

	for (int i = 0; i < PARAM_COUNT; ++i) {
		st[1] += gX[i*3+0] * gX[i*3+0] + gX[i*3+1] * gX[i*3+1] + gX[i*3+2] * gX[i*3+2];
		st[2] += gX[i*3+0] * gY[i*3+0] + gX[i*3+1] * gY[i*3+1] + gX[i*3+2] * gY[i*3+2];
		st[3] += gX[i*3+0] * gZ[i*3+0] + gX[i*3+1] * gZ[i*3+1] + gX[i*3+2] * gZ[i*3+2];
		st[4] += gY[i*3+0] * gY[i*3+0] + gY[i*3+1] * gY[i*3+1] + gY[i*3+2] * gY[i*3+2];
		st[5] += gY[i*3+0] * gZ[i*3+0] + gY[i*3+1] * gZ[i*3+1] + gY[i*3+2] * gZ[i*3+2];
		st[6] += gZ[i*3+0] * gZ[i*3+0] + gZ[i*3+1] * gZ[i*3+1] + gZ[i*3+2] * gZ[i*3+2];
	}


    delete[] gX;
    delete[] gY;
    delete[] gZ;
}

/*************************************************/
/* fibers: fibers that go through that region 	 */
/* n: number of fibers							 */
/* the first fiber in fibers is the one goes	 */ 
/* through the center. 							 */
/*************************************************/
void StructureTensor(CFiber **fibers, const int n, float *st, const int stepSize)
{
	CPoint3F *parameters = new CPoint3F[PARAM_COUNT*n];
	CPoint3F *x = new CPoint3F[n];
	// do the parameterization for each fiber
	for (int i = 0; i < n; ++i) {
		FiberParameterization(fibers[i], fibers[0], &(parameters[i*PARAM_COUNT]), x[i], stepSize);
	}

	
	// estimate the functuon
	float *function = new float[PARAM_COUNT*3*3];
	memset(function, 0, sizeof(float)*PARAM_COUNT*3*3);
	EstimateFunction(parameters, x, n, PARAM_COUNT, function);
	
	// compute the structure tensor
    ComputeStructureTensor(function, st);
    //printf("strength: %f\n", strength);
	
    delete[] parameters;
    delete[] x;
    delete[] function;
}
